import mne
import numpy as np
from scipy.signal import stft
import matplotlib.pyplot as plt

# 文件列表：包含第 3、4、7、8、11 和 12 轮的 EDF 文件
subject = 'S020'
end_subject = '19'
file_paths = [
    fr"D:\下载\eeg-motor-movementimagery-dataset-1.0.0\files\{subject}\{subject}R03.edf",
    fr"D:\下载\eeg-motor-movementimagery-dataset-1.0.0\files\{subject}\{subject}R04.edf",
    fr"D:\下载\eeg-motor-movementimagery-dataset-1.0.0\files\{subject}\{subject}R07.edf",
    fr"D:\下载\eeg-motor-movementimagery-dataset-1.0.0\files\{subject}\{subject}R08.edf",
    fr"D:\下载\eeg-motor-movementimagery-dataset-1.0.0\files\{subject}\{subject}R11.edf",
    fr"D:\下载\eeg-motor-movementimagery-dataset-1.0.0\files\{subject}\{subject}R12.edf",
]

# 定义时间窗口（相对于事件标注的时间，单位：秒）
tmin, tmax = 0.5, 2.5  # 提取事件后 0.5 秒至 2.5 秒之间的数据
# 存储提取的数据
left_data = []
right_data = []

# 定义事件映射
event_id = {'T1': 1, 'T2': 2}

for file_path in file_paths:
    # 读取 EDF 文件（预加载数据）
    raw = mne.io.read_raw_edf(file_path, preload=True, verbose=False)
    raw.filter(l_freq=1., h_freq=40., fir_design='firwin')
    raw.resample(sfreq=128)

    # 从注释中提取事件
    events, _ = mne.events_from_annotations(raw, event_id=event_id)

    # 创建时间片段
    epochs = mne.Epochs(raw, events, event_id=event_id, tmin=tmin, tmax=tmax, baseline=None, preload=True)

    # 提取 T1 和 T2 事件的数据
    if 'T1' in epochs.event_id:
        left_data.append(epochs['T1'].get_data())
    if 'T2' in epochs.event_id:
        right_data.append(epochs['T2'].get_data())

# 将列表转换为 NumPy 数组，方便后续处理
left_data = np.concatenate(left_data, axis=0) if left_data else np.array([])
right_data = np.concatenate(right_data, axis=0) if right_data else np.array([])

# 输出提取的数据形状
print(f"T1（左拳想象）数据形状：{left_data.shape}")
print(f"T2（右拳想象）数据形状：{right_data.shape}")


#   输入：
#       data: EEG trial 数据，形状为 (n_trials, n_channels, n_samples)
#   输出：
#       psd_tensor: 每个通道每个试次的 PSD 特征，形状为 (n_trials, n_channels, n_freq, n_time)
#
def compute_psd_tensor(data, fs, nperseg, noverlap):

    n_trials, n_channels, n_samples = data.shape
    for ch in range(n_channels):
        f, t, Zxx = stft(data[0, ch, :], fs=fs, nperseg=nperseg, noverlap=noverlap)
        break  # 获取频率/时间尺寸
    n_freq = len(f)
    n_time = len(t)

    psd_tensor = np.zeros((n_trials, n_channels, n_freq, n_time))

    for trial in range(n_trials):
        for ch in range(n_channels):
            _, _, Zxx = stft(data[trial, ch, :], fs=fs, nperseg=nperseg, noverlap=noverlap)
            psd = np.abs(Zxx) ** 2  # 功率谱密度
            psd_tensor[trial, ch, :, :] = psd
    return psd_tensor

# 定义 STFT 参数
fs = 250  # 采样频率
nperseg = 64  # 每个段的长度
noverlap = nperseg // 2  # 重叠部分的长度

# 对 T1 和 T2 数据分别进行 STFT 并可视化
if left_data.size > 0:
     left_data_PSD = compute_psd_tensor(left_data, fs, nperseg, noverlap)
     N_left = left_data_PSD.shape[0]
     X_left = np.full(N_left, end_subject, dtype=int)
     Y_left = np.zeros(N_left, dtype=int)
if right_data.size > 0:
     right_data_PSD = compute_psd_tensor(right_data, fs, nperseg, noverlap)
     N_right = right_data_PSD.shape[0]
     X_right = np.full(N_right, end_subject, dtype=int)
     Y_right = np.ones(N_right, dtype=int)

data_identity = np.concatenate((X_left, X_right), axis=0)
data_task = np.concatenate((Y_left, Y_right), axis=0)
data_PSD_all = np.concatenate((left_data_PSD, right_data_PSD), axis=0)
print(data_PSD_all.shape)
data = []
data.append((data_PSD_all, data_identity, data_task))

np.savez(f'D:\pythonproject\DAGN\DATA\session_{end_subject}.npz',
         PSD=data_PSD_all, identity=data_identity, task=data_task)

